// Physical memory allocator, for user processes,
// kernel stacks, page-table pages,
// and pipe buffers. Allocates whole 4096-byte pages.

#include "types.h"
#include "param.h"
#include "memlayout.h"
#include "spinlock.h"
#include "riscv.h"
#include "defs.h"


#define REF_INDEX(a) ((uint64)a/PGSIZE)
#define PAGES_NUM (PHYSTOP / PGSIZE)

void freerange(void *pa_start, void *pa_end);

extern char end[]; // first address after kernel.
                   // defined by kernel.ld.

struct run {
  struct run *next;
};

struct {
  struct spinlock lock;
  struct run *freelist;
} kmem;

// 全局引用结构ref
struct {
  struct spinlock lock; // 对count进行操作的锁
  uint8 count[PAGES_NUM]; // 每个物理页的引用计数
} ref;

void
kinit()
{
  initlock(&kmem.lock, "kmem");
  initlock(&ref.lock, "kmem_ref"); // 初始化锁
  set_ref(); // 初始化引用计数为1
  freerange(end, (void*)PHYSTOP);
}

void
freerange(void *pa_start, void *pa_end)
{
  char *p;
  p = (char*)PGROUNDUP((uint64)pa_start);
  for(; p + PGSIZE <= (char*)pa_end; p += PGSIZE){
    kfree(p);
  }
}

// Free the page of physical memory pointed at by v,
// which normally should have been returned by a
// call to kalloc().  (The exception is when
// initializing the allocator; see kinit above.)
void
kfree(void *pa)
{
  struct run *r;
  if(((uint64)pa % PGSIZE) != 0 || (char*)pa < end || (uint64)pa >= PHYSTOP)
    panic("kfree");

  // free一次，ref计数减1 只有当ref为0，才释放
  acquire(&ref.lock);
  ref.count[REF_INDEX(pa)] -= 1;
  if(get_ref(pa)==0){
    release(&ref.lock); // 保证减完后不不被改变，判断完才能释放锁
    // Fill with junk to catch dangling refs.
    memset(pa, 1, PGSIZE);

    r = (struct run*)pa;

    acquire(&kmem.lock);
    r->next = kmem.freelist;
    kmem.freelist = r;
    release(&kmem.lock);
  }
  else
    release(&ref.lock);
}

// Allocate one 4096-byte page of physical memory.
// Returns a pointer that the kernel can use.
// Returns 0 if the memory cannot be allocated.
void *
kalloc(void)
{
  struct run *r;

  acquire(&kmem.lock);
  r = kmem.freelist;
  // r为待分配和返回的物理页面的地址
  if(r){
    kmem.freelist = r->next;
    acquire(&ref.lock);
    ref.count[REF_INDEX(r)] = 1;
    release(&ref.lock);
  }
  release(&kmem.lock);

  if(r)
    memset((char*)r, 5, PGSIZE); // fill with junk
  return (void*)r;
}


/* Copy on Write Functions */
void set_ref(){
  int i;
  acquire(&ref.lock); 
  for(i=0; i<PAGES_NUM; i++){
    ref.count[i] = 1;
  }
  release(&ref.lock);
}

// 获取引用ref的值
uint8 get_ref(void* pa){
  return ref.count[REF_INDEX(pa)];
}

// 增加引用ref的值
void add_ref(void* pa){
  acquire(&ref.lock); 
  ref.count[REF_INDEX(pa)] += 1;
  release(&ref.lock);
}

// 判断是否是cow页面
int is_cowpage(pagetable_t pagetable, uint64 va) {
  if(va >= MAXVA)
    return 0;
  pte_t* pte = walk(pagetable, va, 0);
  if(pte == 0)
    return 0;
  if((*pte & PTE_V) == 0)
    return 0;
    
  if(*pte & PTE_C) return 1;
  else return 0;
}


// cow分配逻辑
// 发生page fault时，需要为va重新分配物理地址
// 非0: 分配成功 0：分配失败
void* cow_kalloc(pagetable_t pagetable, uint64 va){
  // 将va向下对齐，为其分配整个页
  va = PGROUNDDOWN(va);

  // 获取对应pte
  pte_t* pte = walk(pagetable, va, 0);

  // 获取va对应的物理地址
  uint64 pa = walkaddr(pagetable, va);

  if(get_ref((void *)pa) == 1){
    // 该页面只剩一个引用，直接使用该页面即可
    
    *pte |= PTE_W; // 允许可写
    *pte &= ~PTE_C; // 清除cow标记
    return (void*)pa;
  }
  else{
    // 该页面剩余多个引用，需要重新分配并映射

    // 新分配一个物理页
    char* new_pa = kalloc();
    if(new_pa == 0)
      return 0;
    
    // 重新拷贝一下
    memmove(new_pa, (char*)pa, PGSIZE);

    // 重新添加映射
    *pte &= ~PTE_V; // 标记原映射失效
    *pte |= PTE_W; // 允许可写
    *pte &= ~PTE_C; // 清除cow标记
    if(mappages(pagetable, va, PGSIZE, (uint64)new_pa,PTE_FLAGS(*pte)) != 0) {
      kfree(new_pa);
      *pte |= PTE_V;
      return 0;
    }
    kfree((char*)PGROUNDDOWN(pa));
    return new_pa;
  }


}